import datetime
import gradio as gr
import pymongo as pm
from util import uuid
import time
from http.cookies import SimpleCookie
import redis
import os
import threading
from langchain.llms import OpenAI
from neo4j import GraphDatabase

from PyCmpltrtok.common import sep, to_decode, encode_set, decode_set, list2str_with_limit
from PyCmpltrtok.util_mongo import enqueue, get_history, delete_many_by_user
from common import MONGODB_NAME, VALUE, KEY, IO_PREFIX, CHAT_ONLY, CHAT_AND_DIAG, CHAT_MODES, RESULT_PREFIX, PROMPT_MED
from fastapi_interface import text_gen
from xlangchain_gen import analyze_symptom4user
from python_nlp.xfaiss.flask_faiss_on_mongo_test import search
from python_nlp.kg.neo4j.seek_data_with_embed.xdiagnose import get_symptom, diagnose

os.environ['R_HOST'] = '127.0.0.1'
os.environ['R_PORT'] = '6379'
os.environ['R_PWD'] = 'lgdz4qEdt/ezElyQnXFYXB80iM3OxEbAWRjMFPcIXH5ni6eQ8QOlfp7G7gvV1svPu2Bv7v'
os.environ['M_HOST'] = '127.0.0.1'
os.environ['M_PORT'] = '27017'
os.environ['M_USER'] = 'root'
os.environ['M_PWD'] = 'p1983mdA1-ei'
os.environ['O_PORT'] = '30000'  # openai api port
os.environ['N_PORT'] = '7683'  # neo4j port
os.environ['S_PORT'] = '7700'  # FAISS search service port

S_PORT = int(os.environ['S_PORT'])
DIS_THRESH = 20.0  # 搜索时距离上限
WORDS_LIMIT = 4  # 搜索后最多用这么多症状
THRESH_S_DIAG = 45  # 诊断超过这么多秒就不作数了
CHAR_LIMIT = 450

model = OpenAI(
    streaming=True,
    verbose=True,
    callbacks=[],
    openai_api_key="token1",
    openai_api_base=f"http://127.0.0.1:{os.environ['O_PORT']}/v1",
    model_name="Qwen-1_8B-Chat",
    temperature=0.0,
    openai_proxy=None,
    top_p=1.0,
)

# 连接neo4j
print('Connecting neo4j ...', flush=True)
driver = GraphDatabase.driver(**{
    "uri": f"bolt://127.0.0.1:{os.environ['N_PORT']}",
    "auth": ("neo4j", "p1983ncbdA1-ei"),
    "encrypted": False
})
print('Connected to neo4j.', flush=True)

# 连接Redis
sep('Redis')
rdb = redis.Redis(
    os.environ['R_HOST'], int(os.environ['R_PORT']), 0, socket_timeout=3,
    password=os.environ['R_PWD'],
)
rdb.get('try_it')
sep('Redis OK')

# 连接Mongodb
sep('MongoDB')
mongo = pm.MongoClient(
    os.environ['M_HOST'], int(os.environ['M_PORT']), serverSelectionTimeoutMS=3000,
    username=os.environ['M_USER'], password=os.environ['M_PWD'], 
)
mdb = mongo[MONGODB_NAME]
get_history(mdb, 'user_xxxx', limit=1)  # try it
sep('MongoDB OK')


def add_text(xinput, username):
    if username is None:
        # 为了健壮性
        return '', [], xinput
    
    # 获取聊天历史
    xlog = get_history(mdb, username, more_info=True)
    xlog.append([xinput, None])
    return '', xlog, xinput


def response(xinput, username):
    """输入文本，用模型返回输出。"""    
    if username is None:
        # 为了健壮性
        return []
    
    # 获取聊天历史
    xlog = get_history(mdb, username, more_info=True)
    xlog.append([])
    
    rdb.sadd('users_now', username)
    
    mode = rdb.hget('user2mode', username)
    mode = to_decode(mode, default=CHAT_ONLY)
    
    if CHAT_AND_DIAG != mode:
        # 得到输出
        generator = text_gen(xinput, username)
        for x in generator:
            xlog.pop(-1)
            xlog.append([xinput, x + "▌"])
            yield xlog
            
        # 输入存入mongodb
        ts = time.time_ns()
        enqueue(mdb, 'dialog_in', username, ts, {
            VALUE: xinput,
            'io': 'i',
        })
        
        # 输出存入mongodb
        ts = time.time_ns()
        enqueue(mdb, 'dialog_out', username, ts, {
            VALUE: x,
            'io': 'o',
        })

        xlog.pop(-1)
        xlog.append([xinput, x])
        yield xlog
        rdb.srem('users_now', username)
        return
    else:
        gen = text_gen(xinput, username)
    
        def generate_text():
            rdb.hdel('username2output', username)
            output = ''
            for x in gen:
                rdb.hset('username2dynamic', username, x.encode('utf8'))
                output = x
            rdb.hset('username2output', username, output.encode('utf8'))
                
        generate_text_th = threading.Thread(target=generate_text)
        generate_text_th.start()
        
        # 当前用户的结果
        xencoded = rdb.hget('user2dis', username)
        xencoded = to_decode(xencoded)
        xdis_set = set(decode_set(xencoded))
        xdis_list_ = sorted(xdis_set)
        sep()
        print(f'当前诊断：{xdis_list_}', flush=True)
        
        xsymptoms = analyze_symptom4user(mdb, xinput, username, model)
        
        # 搜索输入，选出距离小于阈值的结果
        for xsym in xsymptoms:
            sep(xsym)
            xoutput = search(xsym, 8, S_PORT)
            xsents = xoutput['sentences']
            xds = xoutput['D']
            print('候选:', flush=True)
            xwords = []
            for i, s in enumerate(xsents):
                d = xds[i]
                print(i, s, d, d < DIS_THRESH, flush=True)
                if d < DIS_THRESH:
                    xwords.append(s)
        xwords = sorted(set(xwords))

        # 从这些结果中挑选出是“症状”的节点名字
        xlist = get_symptom(driver, xwords, WORDS_LIMIT)
        print(f'输入了症状：{xlist}', flush=True)
        
        # 用这些症状去诊断
        xdiseases = diagnose(driver, xlist)
        if xdiseases:
            print(f'本轮诊断：{xdiseases}', flush=True)
            if not xdis_set:
                # 第一次诊断
                xdis_set = set(xdiseases)
            else:
                # 后续诊断
                xresult = xdis_set.intersection(set(xdiseases))  # 求交集
                if len(xresult):
                    print('【有交集】', flush=True)
                else:
                    # 没有交集求并集
                    print('【没有交集求并集】', flush=True)
                    xresult = xdis_set.union(set(xdiseases))
                xdis_set = xresult
                print('交并后:', xdis_set, flush=True)
                    
        # 结果            
        xdis_list = sorted(xdis_set)
        
        # 结果还是那样，就不使用
        if xdis_list_ == xdis_list:
            print('【诊断结果没有变化】', flush=True)
            while True:
                xoutput = rdb.hget('username2output', username)
                xoutput = to_decode(xoutput, default=None)
                
                if xoutput is not None:
                    xlog.pop(-1)
                    xlog.append([xinput, xoutput])
                    yield xlog
                    
                    # 输入存入mongodb
                    ts = time.time_ns()
                    enqueue(mdb, 'dialog_in', username, ts, {
                        VALUE: xinput,
                        'io': 'i',
                    })
                    
                    # 输出存入mongodb
                    ts = time.time_ns()
                    enqueue(mdb, 'dialog_out', username, ts, {
                        VALUE: xoutput,
                        'io': 'o',
                    })
                    
                    return
                else:
                    x = rdb.hget('username2dynamic', username)
                    x = to_decode(x)
                    xlog.pop(-1)
                    xlog.append([xinput, x + "▌"])
                    yield xlog    
                
        
        # 保存用户结果
        xencoded = encode_set(xdis_set)
        rdb.hset('user2dis', username, xencoded.encode('utf8'))
        
        # 返回诊断
        xoutput = RESULT_PREFIX + list2str_with_limit(xdis_list, CHAR_LIMIT)
        if len(xdis_list) > 5:
            xoutput += '（可能性较多，请进一步描述症状来缩小范围。）'
        xoutput += PROMPT_MED
        xlog.pop(-1)
        xlog.append([xinput, xoutput])
        yield xlog
        
        # 输入存入mongodb
        ts = time.time_ns()
        enqueue(mdb, 'dialog_in', username, ts, {
            VALUE: xinput,
            'io': 'i',
        })
        
        # 输出存入mongodb
        ts = time.time_ns()
        enqueue(mdb, 'dialog_out', username, ts, {
            VALUE: xoutput,
            'io': 'o',
        })
        

def embed_user_id_in_state_comp(request: gr.Request):
    """
    页面加载的事件handler

    从HTTP request中获取cookie来确定用户名，如果没有则生成一个
    :param request: HTTP request data
    :return: 将用户名发送给浏览器中存放用户名的组件
    """
    xuuid = uuid()

    #从cookie中获取username
    username = None
    try:
        cookie_str = request.headers.cookie
        print('cookie_str:', cookie_str, flush=True)
        cookie = SimpleCookie()
        cookie.load(cookie_str)
        username = cookie.get('username', None)
        print('user:', username, flush=True)
        username = username.value
    except Exception:
        try:
            username = request.cookies['username']
        except Exception:
            pass

    # 获取不到，新建一个username
    if username is None:
        print('new user!', flush=True)
        username = f'u{xuuid}'

    # print username
    print('user:', username, flush=True)

    xlog = get_history(mdb, username, more_info=True)

    return username, username, username, xlog


# 用户名写入cookie的JS
js_username2cookie = """
    async () => {{
            setTimeout(() => {{
                var xnew = document.querySelector('#username.prose').innerHTML;
                console.log('username', xnew)
                if ('' == xnew) {
                    return;
                }
                document.cookie = 'username=' + xnew + '; max-age=' + 3600*24*30*6;
            }}, 1);
        }}
"""


def do_the_init(username):
    """
        点击”开始聊天“的事件的handler
        :return: 隐藏”开始聊天“，显示其他组件。
    """
    if username is None:
        # 为了健壮性
        mode = CHAT_ONLY
    else:
        mode = rdb.hget('user2mode', username)
        mode = to_decode(mode, default=CHAT_ONLY)
    return (
        mode,
        gr.update(visible=False),
        *([gr.update(visible=True)] * 3),
    )
    
    
def show_active_user_count():
    users_count = rdb.scard('users_now')
    text = f'【活跃用户{users_count}人】'
    return text


def clear_history(username):
    if username is None:
        # 为了健壮性
        return []
    delete_many_by_user(mdb, 'dialog_in', username)
    delete_many_by_user(mdb, 'dialog_out', username)
    return []


def change_mode(username, mode):
    if username is None:
        # 为了健壮性
        return CHAT_ONLY
    rdb.hset('user2mode', username, mode.encode('utf8'))
    return mode
    

# 构建Gradio界面
with gr.Blocks(analytics_enabled=False) as demo:

    # 用户名相关控件
    cmp_username_state = gr.State()
    cmp_username_html = gr.HTML(visible=False, elem_id='username')
    cmp_username_display = gr.Textbox(interactive=False, label='用户名')
    
    cmp_users_count_display = gr.Textbox(interactive=False, label='活跃用户')

    # 启动按钮
    cmp_start_btn = gr.Button('开始聊天')

    # 聊天历史
    cmp_chatbot = gr.Chatbot().style(height=350)
    
    # 输入文本框
    cmp_input = gr.Textbox(interactive=True, label='输入', visible=False, )

    cmp_input_state = gr.State()

    # 输入按钮
    cmp_input_button = gr.Button('发送', visible=False, )

    cmp_choose_mode = gr.Radio([CHAT_ONLY, CHAT_AND_DIAG], value=CHAT_ONLY, label='选择对话模式')
    cmp_choose_mode.visible = False

    cmp_clear_button = gr.Button('清理历史')

    cmp_clear_button.click(clear_history, cmp_username_state, cmp_chatbot, queue=False)

    # 输入文本框回车
    cmp_input.submit(add_text, [cmp_input, cmp_username_state], [
        cmp_input,
        cmp_chatbot,
        cmp_input_state,
    ], queue=False).then(response, [cmp_input_state, cmp_username_state], cmp_chatbot)

    # 输入按钮点击事件
    cmp_input_button.click(add_text, [cmp_input, cmp_username_state], [
        cmp_input,
        cmp_chatbot,
        cmp_input_state,
    ], queue=False).then(response, [cmp_input_state, cmp_username_state], cmp_chatbot)

    cmp_choose_mode.change(change_mode, [cmp_username_state, cmp_choose_mode], cmp_choose_mode)

    # 加载时获取或设置cookie，并展示聊天历史
    demo.load(embed_user_id_in_state_comp, None, [
        cmp_username_state,
        cmp_username_display,
        cmp_username_html,
        cmp_chatbot,
    ], queue=False)
    # cmp_username_html发生变化时，触发JS，把生成的username放入页面cookie
    cmp_username_html.change(None, None, None, _js=js_username2cookie, queue=False)
    
    # 启动
    cmp_start_btn.click(do_the_init, cmp_username_state, [
        cmp_choose_mode,
        cmp_start_btn,
        cmp_input,
        cmp_input_button,
        cmp_choose_mode,
    ], queue=False)
    
    cmp_start_btn.click(show_active_user_count, None, cmp_users_count_display, every=0.2)
    
# 使用queue（目前这个例子太简单，没有用到queue，后面会用到。）
demo.queue(
    concurrency_count=10,
    status_update_rate='auto',
    # status_update_rate=0.02,
)

# 启动Gradio
# demo.launch(server_name='0.0.0.0', server_port=6006, share=True, debug=False)  # 带内网穿透
demo.launch(server_name='0.0.0.0', server_port=6006, debug=False, inbrowser=False)
